#define K_MAX_SHAPE_DIM 0

#include "kernel_operator.h"
using namespace AscendC;

class KernelEye {
public:
    __aicore__ inline KernelEye() {}
    __aicore__ inline void Init(GM_ADDR y, uint16_t num_columns,
                                uint16_t num_rows, uint32_t batchNum, TPipe* pipeIn) {
        this->num_columns = num_columns;
        this->num_rows = num_rows;
        this->batchNum = batchNum;
        this->batchSize = num_rows*num_columns;
        this->pipe = pipeIn;
        uint32_t coreNum = GetBlockNum();
        uint32_t coreIdx = GetBlockIdx();

        uint32_t tail = batchNum%coreNum;
        uint32_t coreDataNum = batchNum/coreNum;
        // yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, 10);
        uint32_t bigDataNum = coreDataNum + 1;
        uint32_t globalIndex = bigDataNum * coreIdx;
        if (coreIdx < tail) {
            this->coreDataNum = bigDataNum;
        } else {
            this->coreDataNum = coreDataNum;
            globalIndex -= (coreIdx - tail);
        }
        loop = (num_rows < num_columns) ? num_rows:num_columns;
        if (loop > num_columns) loop = num_columns;
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y + globalIndex*batchSize, this->coreDataNum*batchSize *sizeof(DTYPE_Y));
        pipe->InitBuffer(tmpBuf, 32*loop);
        // pipe->InitBuffer(tmpOut,1, 32*loop);

    }
    __aicore__ inline void process() {
        LocalTensor<DTYPE_Y> temp = tmpBuf.Get<DTYPE_Y>();
        uint32_t blockLen = 32/sizeof(DTYPE_Y);
        // LocalTensor<DTYPE_Y> temp = tmpOut.AllocTensor<DTYPE_Y>();
        if constexpr (std::is_same_v<DTYPE_Y, double>){
            for (int i = 0; i < loop; i++) {
                temp.SetValue(i*blockLen, DTYPE_Y(1));
            }
        } else {
            Duplicate(temp, DTYPE_Y(1), blockLen*loop);
        }
        // int32_t index, t;
        // int32_t batch_num = num_columns*num_rows;
        AscendC::DataCopyExtParams copyParams{loop, sizeof(DTYPE_Y), 0, static_cast<uint32_t>(num_columns*sizeof(DTYPE_Y)), 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位

        for (int i = 0; i < coreDataNum*batchSize; i+=batchSize) {
            DataCopyPad(yGm[i],temp ,copyParams);
            // for(int32_t j = 0; j < loop; j++){
            //     // if(j < num_columns){
            //         t = j * num_columns + j;
            //         index = i * batchSize + t;
            //         yGm.SetValue(index, 1);
            // }
        }
    }
private:
    TPipe* pipe;
    GlobalTensor<DTYPE_Y> yGm;
    TBuf<QuePosition::VECCALC> tmpBuf;
    // TQue<QuePosition::VECOUT,1>tmpOut;
    uint16_t loop;
    uint16_t num_columns;
    uint16_t num_rows;
    uint32_t batchNum;
    uint32_t batchSize;
    uint32_t coreDataNum;

};
extern "C" __global__ __aicore__ void eye(GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelEye op;
    TPipe pipe;

    op.Init(y, tiling_data.num_columns, tiling_data.num_rows,tiling_data.batchNum,&pipe);
    op.process();
}